from enum import Enum
import math
import time
import numpy as np
import serial.tools.list_ports
import serial
import sys
import pyqtgraph as pg
from pyqtgraph.Qt import QtWidgets
import threading

class GSRApp:
    command_end = "\r\n"
    class Command(Enum):
        HARDWARE_VERSION = "AT+HV"
        SOFTWARE_VERSION = "AT+SV"
        BAUD_RATE = "AT+BD"
        MODE = "AT+MD"
        RESISTANC_GRADE = "AT+RG"
        UNIT_TYPE = "AT+UT"

    unit_types = ["kΩ", "nA"]
    grade_gain = [1000, 100, 10, 1]
    
    def __init__(self):
        self.app = pg.mkQApp("Plotting Example")
        self.win = QtWidgets.QMainWindow()
        self.win.setWindowTitle('GSR生物电阻')
         # Central widget
        central_widget = QtWidgets.QWidget()
        self.win.setCentralWidget(central_widget)

        # Layout
        layout = QtWidgets.QVBoxLayout()
        central_widget.setLayout(layout)

        # First row: Connect button and Status label
        toolbar_layout = QtWidgets.QHBoxLayout()

        self.connect_btn = QtWidgets.QPushButton("Connect")
        self.connect_btn.setMaximumWidth(100)
        self.connect_btn.clicked.connect(self.toggle_connection)
        toolbar_layout.addWidget(self.connect_btn)

        toolbar_layout.addWidget(QtWidgets.QLabel("Baud Rate:"))
        self.baud_rate_combobox = QtWidgets.QComboBox()
        self.baud_rate_combobox.addItems([str(rate) for rate in [300, 1200, 2400, 4800, 9600, 19200, 38400, 57600, 115200]])
        self.baud_rate_combobox.setCurrentIndex(self.baud_rate_combobox.findText("9600"))
        toolbar_layout.addWidget(self.baud_rate_combobox)

        toolbar_layout.addWidget(QtWidgets.QLabel("Resistance Grade:"))
        self.resistance_grade_combobox = QtWidgets.QComboBox()
        self.resistance_grade_combobox.addItems(['1k~100k', '10k~1000k', '100k~10000k', '1000k~100000k'])
        self.resistance_grade_combobox.setCurrentIndex(0)
        self.resistance_grade_combobox.currentIndexChanged.connect(lambda: self.send_command(self.Command.RESISTANC_GRADE, self.resistance_grade_combobox.currentIndex()+1))
        toolbar_layout.addWidget(self.resistance_grade_combobox)

        self.status_label = QtWidgets.QLabel("Not Connected")
        toolbar_layout.addWidget(self.status_label)

        self.frequency_label = QtWidgets.QLabel("Frequency: 0.0 Hz")
        toolbar_layout.addWidget(self.frequency_label)

        self.realtime_resistance_label = QtWidgets.QLabel("Realtime Resistance: 0.0 kΩ")
        toolbar_layout.addWidget(self.realtime_resistance_label)

        self.clear_btn = QtWidgets.QPushButton("Clear")
        self.clear_btn.setMaximumWidth(100)
        self.clear_btn.clicked.connect(self.clear_data)
        toolbar_layout.addWidget(self.clear_btn)

        layout.addLayout(toolbar_layout)
        # Second row: Plot
        self.signal_plot = pg.PlotWidget(title="GSR")
        self.data_curve = self.signal_plot.plot(pen='w')
        self.moving_avg_curve = self.signal_plot.plot(pen='y')
        self.local_max_curve = self.signal_plot.plot(symbolBrush='g')
        self.local_min_curve = self.signal_plot.plot(symbolBrush='b')
        layout.addWidget(self.signal_plot)


        self.serial = None
        self.is_listening = False
        # Data storage
        self.timestamps = []
        self.data = []
        self.moving_averages = []
        self.local_max = []
        self.local_min = []
        self.last_update_time = 0
        self.command_queue= []
        self.app_lock = threading.Lock()
        self.resistance_grade = 1
        self.unit_type = 0
        self.frequency = 1
        self.start_timestamp = None

    def run(self):
        self.win.showMaximized()
        self.app.exec()

    def send_command(self, command, value=None):
        self.app_lock.acquire()
        if value is None:
            self.command_queue.append(command)
        else:
            self.command_queue.append((command, value))
        self.app_lock.release()
    def toggle_connection(self):
        if self.is_listening:
            self.is_listening = False
        else:
            ports = serial.tools.list_ports.comports()
            available_ports = [port.device for port in ports]
            if len(available_ports) == 0:
                self.serial = None
            else:
                try:
                    port = available_ports[0]                    
                    self.status_label.setText(f"Connecting to {available_ports[0]}...")
                    # Connect to the device (e.g., COM1 at 57600 baud)
                    baud_rate = int(self.baud_rate_combobox.currentText())
                    self.serial = serial.Serial(port, baud_rate, bytesize=8, stopbits=1, parity='N')
                    if self.serial.is_open:
                        self.status_label.setText(f"{available_ports[0]} connected with baud rate {self.serial.baudrate}")
                        self.connect_btn.setText("Disconnect")
                        self.send_command(self.Command.HARDWARE_VERSION)
                        self.send_command(self.Command.SOFTWARE_VERSION)
                        self.send_command(self.Command.BAUD_RATE)
                        self.send_command(self.Command.MODE, 0)
                        self.send_command(self.Command.RESISTANC_GRADE, self.resistance_grade_combobox.currentIndex()+1)
                        self.send_command(self.Command.UNIT_TYPE)
                        self.start_timestamp = time.time()
                        threading.Thread(target=self.listen, daemon=True).start()
                except serial.SerialException as e:
                    self.status_label.setText(f"连接串口 {port} 时出错: {e}")
                    self.serial = None
                
    def run_command(self, cmd):
        if isinstance(cmd, tuple):
            cmd, value = cmd[0], cmd[1]
            cmd_str = f"{cmd.value}:{value}{self.command_end}"
        else:
            value = None
            cmd_str = f"{cmd.value}{self.command_end}"

        print(f'[SEND]{cmd_str.strip('\r\n')}')
        self.serial.write(cmd_str.encode('utf-8'))
        self.serial.flush()

        while True:
            ret = self.serial.readline().decode('utf-8').strip(' \r\n')
            print(f'[RECV]{ret}')
            if ret.startswith('U'):
                continue

            match cmd:
                case self.Command.HARDWARE_VERSION:
                    if ret.startswith('RES '):
                        hardware_version = ret.split(' ')[-1]
                        print(f'hardware_version: {hardware_version}\n')
                        break
                case self.Command.SOFTWARE_VERSION:
                    if ret.startswith('RES '):
                        software_version = ret.split(' ')[-1]
                        print(f'software_version: {software_version}\n')
                        break                
                case _:
                    if ret.startswith(cmd.value):
                        print('\n')
                        match cmd:
                            case self.Command.RESISTANC_GRADE:
                                self.resistance_grade = int(ret.split(':')[-1])
                                break
                            case self.Command.UNIT_TYPE:
                                self.unit_type = int(ret.split(':')[-1])
                                break
                        break
    def listen(self):
        if self.is_listening:
            return
        
        self.is_listening = True

        # Continuously read and process data
        while self.is_listening:
            if len(self.command_queue) > 0:
                self.app_lock.acquire()
                for cmd in self.command_queue:                                       
                    self.run_command(cmd)
                
                self.command_queue = []
                self.app_lock.release()

            ret = self.serial.readline().decode('utf-8').strip(' \r\n')
            print(f'[RECV]{ret}')                 
            if ret.startswith("U"):
                timestamp = time.time()
                mode = int(ret[1])
                data = int(ret[3:])
                unit = self.unit_types[self.unit_type]
                # print(f'[{ret}] {data}{unit}')                
                self.update_data(timestamp, data)

                    
        self.serial.close()
        self.serial = None
        self.connect_btn.setText("Connect")
        self.status_label.setText("Not Connected")
    def update_data(self, timestamp, value):   
        timestamp -= self.start_timestamp     
        self.timestamps.append(timestamp)     
        gain = self.grade_gain[self.resistance_grade - 1]
        value /= gain
        self.data.append(value)
        self.update_local_attrs(window=1)
        if timestamp - self.last_update_time > 0.1:
            self.app_lock.acquire()
            self.last_update_time = timestamp
            if len(self.timestamps) > 5000:
                data_slice = slice(-5000, len(self.timestamps))
            else:
                data_slice = slice(0, len(self.timestamps))
            self.data_curve.setData(self.timestamps[data_slice], self.data[data_slice])
            self.moving_avg_curve.setData(self.timestamps[data_slice], self.moving_averages[data_slice])
            # self.local_voltage_curve.setData(self.timestamps[data_slice], self.local_voltages[data_slice])
            # self.local_max_curve.setData(self.timestamps[self.local_max], self.data[self.local_max])
            # self.local_min_curve.setData(self.timestamps[self.local_min], self.data[self.local_min])
            self.signal_plot.setXRange(self.timestamps[-1] - 10, self.timestamps[-1])

            unit = self.unit_types[self.unit_type]
            self.realtime_resistance_label.setText(f"Realtime Resistance: {value}{unit}")
            if len(self.timestamps) > 10:
                self.frequency = int(10 / (self.timestamps[-1] - self.timestamps[-10]))
                self.frequency_label.setText(f"Frequency: {self.frequency} Hz")

            self.app_lock.release()

    def update_local_attrs(self, window):
        if len(self.timestamps) == len(self.moving_averages):
            return
        
        assert len(self.timestamps) > len(self.moving_averages)

        if len(self.timestamps) == 1:
            self.local_max.append(0)
            self.local_min.append(0)
            self.moving_averages.append(self.data[0])
            return
        
        for last in range(len(self.moving_averages), len(self.timestamps)):            
            last_timestamp = self.timestamps[last]
            find_max = True
            if last_timestamp - self.timestamps[self.local_max[last - 1]] <= window:
                find_max = False
                if self.data[last] >= self.data[self.local_max[last - 1]]:
                    self.local_max.append(last)
                else:
                    self.local_max.append(self.local_max[last - 1])

            find_min = True
            if last_timestamp - self.timestamps[self.local_min[last - 1]] <= window:
                find_min = False
                if self.data[last] <= self.data[self.local_min[last - 1]]:
                    self.local_min.append(last)
                else:
                    self.local_min.append(self.local_min[last - 1])
            
            local_max = local_min = i = len(self.timestamps) - 1
            while i >= 0 and self.timestamps[last] - self.timestamps[i] <= window:
                if find_max and self.data[i] > self.data[local_max]:
                    local_max = i
                elif find_min and self.data[i] < self.data[local_min]:
                    local_min = i
                i -= 1

            if find_max:
                self.local_max.append(local_max)

            if find_min:
                self.local_min.append(local_min)

            self.moving_averages.append((self.moving_averages[-1] * (self.frequency - 1) + self.data[last]) / self.frequency)

    def clear_data(self):
        self.app_lock.acquire()
        self.timestamps = []
        self.data = []
        self.moving_averages = []
        self.local_max = []
        self.local_min = []
        self.data_curve.setData(self.timestamps, self.data)
        self.moving_avg_curve.setData(self.timestamps, self.moving_averages)
        self.app_lock.release()

if __name__ == '__main__':
    app = GSRApp()
    app.run()